#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author   : Uyynot
# @Email    : uyynot@qq.com
# @Time     : 2025/7/17 11:35
# @File     : d.py
# @Project  : demoProject
# @Desc     :
import functools
import logging
import time
from contextvars import ContextVar
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Dict, Any, List

from pydantic import BaseModel
class QueryState(Enum):
    INITIAL = auto()
    ENTITY_RETRIEVAL = auto()
    QA_RETRIEVAL = auto()
    SQL_GENERATION = auto()
    INTENT_RECOGNITION = auto()
    SEARCH_INTENT = auto()
    AGENT_SEARCH = auto()
    REJECT_INTENT = auto()
    KNOWLEDGE_SEARCH = auto()
    EXECUTE_QUERY = auto()
    ANALYZE_DATA = auto()
    AGENT_TASK = auto()
    AGENT_DATA_SUMMARY = auto()
    ASK_ENTITY_SELECT = auto()
    ASK_QUERY_REWRITE = auto()
    QUERY_REWRITE = auto()
    USER_SELECT_ENTITY = auto()
    DATA_VISUALIZATION = auto()
    ERROR = auto()
    COMPLETE = auto()
    QUERY_ENTITY = auto()

@dataclass
class ProcessingContext:
    search_box: str
    query_rewrite: str
    session_id: str
    user_id: str
    username: str
    selected_profile: str
    database_profile: Dict[str, Any]
    model_type: str
    use_rag_flag: bool
    intent_ner_recognition_flag: bool
    agent_cot_flag: bool
    explain_gen_process_flag: bool
    visualize_results_flag: bool
    data_with_analyse: bool
    gen_suggested_question_flag: bool
    auto_correction_flag: bool
    context_window: int
    entity_same_name_select: Dict[str, Any]
    user_query_history: List[str]
    opensearch_info: Dict[str, Any]
    previous_state: str = "INITIAL"
    entity_retrieval: List[str] = field(default_factory=list)
    entity_user_select: Dict[str, Any] = field(default_factory=dict)
    limit: int = 100

class KnowledgeSearchResult(BaseModel):
    knowledge_response: str

class ChartEntity(BaseModel):
    chart_type: str
    chart_data: list[Any]
class SQLSearchResult(BaseModel):
    sql: str
    sql_data: list[Any]
    sql_data_url: str = ""
    data_show_type: str
    sql_gen_process: str
    data_analyse: str
    sql_data_chart: list[ChartEntity]
class TaskSQLSearchResult(BaseModel):
    sub_task_query: str
    sql_search_result: SQLSearchResult
class AgentSearchResult(BaseModel):
    agent_sql_search_result: list[TaskSQLSearchResult]
    agent_summary: str
class AskReplayResult(BaseModel):
    query_rewrite: str

class AskEntitySelect(BaseModel):
    entity_select_info: dict[str, Any]
    entity_retrieval: list[Any]
class Answer(BaseModel):
    data_profiles: str = ""
    query: str
    query_rewrite: str = ""
    query_intent: str
    knowledge_search_result: KnowledgeSearchResult
    sql_search_result: SQLSearchResult
    agent_search_result: AgentSearchResult
    ask_rewrite_result: AskReplayResult
    suggested_question: list[str]
    ask_entity_select: AskEntitySelect
    error_log: dict[str, Any]

LOG_LEVEL='INFO'

trace_id_var: ContextVar[str] = ContextVar('trace_id', default='')
class TraceFormatter(logging.Formatter):
    """自定义日志格式化器，支持trace_id"""

    def format(self, record):
        # 从上下文变量中获取trace_id
        trace_id = trace_id_var.get('')
        record.trace_id = trace_id if trace_id else 'N/A'
        return super().format(record)
def getLogger():
    global logger
    if logger is not None:
        return logger

    # 创建日志记录器
    logger = logging.getLogger('application')
    logger.propagate = False
    logger.setLevel(LOG_LEVEL)
    # 创建控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(LOG_LEVEL)

    # 设置日志格式，包含trace_id
    log_format = '%(asctime)s [%(module)s L%(lineno)d] [%(levelname)s] [trace_id:%(trace_id)s] %(message)s'
    formatter = TraceFormatter(log_format)

    # 设置日志处理器格式
    console_handler.setFormatter(formatter)

    # 清理旧的日志处理器
    logger.handlers.clear()
    # 添加日志处理器
    logger.addHandler(console_handler)

    return logger

logger = getLogger()
def log_execution(func):
    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        state_name = self.get_state().name if self.get_state() else "Unknown State"
        logger.info(f"Executing {func.__name__} in state {state_name}")
        return func(self, *args, **kwargs)

    return wrapper

class QueryStateMachine:
    def __init__(self, context: ProcessingContext):
        self.state = self.get_state_from_name("INITIAL")
        self.previous_state = self.get_state_from_name(context.previous_state)
        self.context = context
        self.answer = Answer(
            query="",
            query_rewrite="",
            query_intent="",
            knowledge_search_result=KnowledgeSearchResult(knowledge_response=""),
            sql_search_result=SQLSearchResult(
                sql="",
                sql_data=[],
                data_show_type="table",
                sql_gen_process="",
                data_analyse="",
                sql_data_chart=[],
            ),
            agent_search_result=AgentSearchResult(
                agent_sql_search_result=[], agent_summary=""
            ),
            ask_rewrite_result=AskReplayResult(query_rewrite=""),
            suggested_question=[],
            ask_entity_select=AskEntitySelect(
                entity_select_info={}, entity_retrieval=[]
            ),
            error_log={},
        )

        self.search_intent_flag = False
        self.reject_intent_flag = False
        self.agent_intent_flag = False
        self.knowledge_search_flag = False

        self.intent_search_result = {}
        self.intent_response = {}
        self.entity_slot = []
        self.normal_search_entity_slot = []
        self.normal_search_qa_retrival = []
        self.agent_cot_retrieve = []
        self.agent_task_split = {}
        self.agent_search_result = []
        self.agent_data_analyse_result = ""
        self.agent_valid_data = []
        self.use_auto_correction_flag = False
        self.first_sql_execute_info = {}
        self.token_info = {}

        self.query_time_replace = ""

    def transition(self, new_state):
        self.state = new_state

    def get_answer(self):
        return self.answer

    def get_state(self):
        return self.state

    def get_state_from_name(self, state_name):
        if state_name == QueryState.INITIAL.name:
            return QueryState.INITIAL
        elif state_name == QueryState.USER_SELECT_ENTITY.name:
            return QueryState.USER_SELECT_ENTITY

    def run(self):
        if self.previous_state == QueryState.USER_SELECT_ENTITY:
            self.transition(QueryState.USER_SELECT_ENTITY)

        while self.state != QueryState.COMPLETE and self.state != QueryState.ERROR:
            if self.state == QueryState.INITIAL:
                self.handle_initial()
            elif self.get_state() == QueryState.REJECT_INTENT:
                self.handle_reject_intent()
            elif self.get_state() == QueryState.KNOWLEDGE_SEARCH:
                self.handle_knowledge_search()
            elif self.state == QueryState.ENTITY_RETRIEVAL:
                self.handle_entity_retrieval()
            elif self.state == QueryState.QA_RETRIEVAL:
                self.handle_qa_retrieval()
            elif self.state == QueryState.SQL_GENERATION:
                self.handle_sql_generation()
            elif self.state == QueryState.INTENT_RECOGNITION:
                self.handle_intent_recognition()
            elif self.state == QueryState.EXECUTE_QUERY:
                self.handle_execute_query()
            elif self.state == QueryState.ANALYZE_DATA:
                self.handle_analyze_data()
            elif self.state == QueryState.ASK_ENTITY_SELECT:
                self.handle_entity_selection()
            elif self.state == QueryState.AGENT_TASK:
                self.handle_agent_task()
            elif self.state == QueryState.AGENT_SEARCH:
                self.handle_agent_sql_generation()
            elif self.state == QueryState.AGENT_DATA_SUMMARY:
                self.handle_agent_analyze_data()
            elif self.state == QueryState.USER_SELECT_ENTITY:
                self.handle_user_select_entity()
            else:
                self.state = QueryState.ERROR

        if self.state == QueryState.COMPLETE:
            self.handle_data_visualization()
        if self.context.gen_suggested_question_flag:
            self.handle_suggest_question()

    @log_execution
    def handle_initial(self):
        """
        处理初始状态，包括查询实体抽取、问题改写和意图识别
        添加耗时日志记录以监控性能
        """
        start_time = time.time()
        logger.info(f"handle_initial_begin 开始执行，查询内容: {self.context.search_box}")

        try:
            self.answer.query = self.context.search_box
            self.answer.query_rewrite = self.context.search_box
            self.answer.query_intent = "normal_search"

            # 使用统一的LLM调用替代原有的三次独立调用
            from utils.llm import get_unified_query_analysis
            logger.info(f"get_unified_query_analysis user_query_history:  {self.context.user_query_history}")
            unified_result, model_response = get_unified_query_analysis(
                self.context.model_type,
                self.context.search_box,
                self.context.database_profile["prompt_map"],
                self.context.user_query_history,
                self.context.database_profile["prompt_environment"]
            )

            # 记录token信息
            self.token_info[QueryState.INITIAL.name] = model_response.token_info

            # 处理统一结果，进行字段映射
            query_entity_result, query_rewrite_result, intent_result = self._process_unified_result(unified_result)
            # 处理实体抽取结果（时间实体替换逻辑）
            self._process_entity_result(query_entity_result)

            # 处理问题改写结果
            if self.context.context_window > 0:
                # 多轮会话场景，判断
                try:
                    query_rewrite_intent = query_rewrite_result.get("intent")
                    self.context.query_rewrite = query_rewrite_result.get("query")
                    logger.info(f"context_window >  0 self.context.query_rewrite = {self.context.query_rewrite}")
                    self.answer.query_rewrite = self.context.query_rewrite
                    # 意图不清晰，继续追问
                    if query_rewrite_intent == "ask_in_reply":
                        self._set_ask_in_reply_result()
                    else:
                        # 直接处理意图分类结果并进行状态转换
                        self._process_intent_result(intent_result)
                except Exception as e:
                    logger.error(f"Error processing rewrite result: {e}")
                    self.context.query_rewrite = self.query_time_replace
                    self.transition(QueryState.INTENT_RECOGNITION)
            else:
                self.context.query_rewrite = self.query_time_replace
                logger.info(f"context_window = 0 self.context.query_rewrite = {self.context.query_rewrite}")
                self.answer.query_rewrite = self.query_time_replace
                # 直接处理意图分类结果并进行状态转换
                self._process_intent_result(intent_result)

        except Exception as e:
            self.answer.error_log[QueryState.INITIAL.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, handle_initial encountered an error: {e}"
            )
            self.transition(QueryState.ERROR)
        finally:
            end_time = time.time()
            execution_time = end_time - start_time
            logger.info(f"handle_initial_end 执行完成，总耗时: {execution_time:.3f}秒")

    @log_execution
    def _handle_query_entity(self):
        try:
            query_entity, model_response = get_query_entity(
                self.context.model_type,
                self.context.search_box,
                self.context.database_profile["prompt_map"],
                self.context.database_profile["prompt_environment"],
            )
            self.token_info[QueryState.QUERY_ENTITY.name] = model_response.token_info
            time_entity = query_entity.get("time_entity", {})
            self.query_time_replace = self.context.search_box
            time_slot_list = []
            if len(time_entity) > 0:
                for key, value in time_entity.items():
                    time_slot_list.append(value)

            time_change_slot_list = DateTimeUtil.run(time_slot_list)
            time_change_slot_dict = {}
            for item in time_change_slot_list:
                entity_item = item.split("=")
                if len(entity_item) == 2:
                    time_change_slot_dict[entity_item[0]] = entity_item[1]
            slot_replace = {}
            if len(time_entity) > 0:
                for key, value in time_entity.items():
                    if value in time_change_slot_dict:
                        slot_replace[key] = time_change_slot_dict[value]

            logger.info(slot_replace)
            for key, value in slot_replace.items():
                self.query_time_replace = self.query_time_replace.replace(key, value, 1)
            logger.info(f"context.search_box: {self.context.search_box}")
            logger.info(f"query_time_replace: {self.query_time_replace}")

        except Exception as e:
            traceback.print_exc()
            self.answer.error_log[QueryState.QUERY_ENTITY.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, handle_initial encountered an error: {e}"
            )
            self.transition(QueryState.ERROR)

    def _process_unified_result(self, unified_result):
        """
        处理统一LLM调用的结果，进行字段映射和兼容性处理

        Args:
            unified_result: 统一LLM调用返回的结果字典

        Returns:
            tuple: (query_entity_result, query_rewrite_result, intent_result)
        """
        try:
            # 提取各部分结果
            rewrite_intent = unified_result.get("rewrite_intent", "original_problem")
            query_rewrite = unified_result.get("query", self.context.search_box)
            classification_intent = unified_result.get("classification_intent", "normal_search")
            normal_entity = unified_result.get("normal_entity", [])
            time_entity = unified_result.get("time_entity", {})

            # 构造query_entity结果（兼容原有格式）
            query_entity_result = {
                "normal_entity": normal_entity,
                "time_entity": time_entity
            }

            # 构造query_rewrite结果（字段映射：rewrite_intent -> intent）
            query_rewrite_result = {
                "intent": rewrite_intent,  # rewrite_intent映射为intent
                "query": query_rewrite
            }

            # 构造intent结果（字段映射：classification_intent -> intent）
            intent_result = {
                "intent": classification_intent  # classification_intent映射为intent
            }

            logger.info(f"Unified result processed - rewrite_intent: {rewrite_intent}, "
                        f"classification_intent: {classification_intent}, "
                        f"normal_entity: {normal_entity}, time_entity: {time_entity}")

            return query_entity_result, query_rewrite_result, intent_result

        except Exception as e:
            logger.error(f"Error processing unified result: {e}")
            # 返回默认值
            default_entity = {"normal_entity": [], "time_entity": {}}
            default_rewrite = {"intent": "original_problem", "query": self.context.search_box}
            default_intent = {"intent": "normal_search"}
            return default_entity, default_rewrite, default_intent

    def _process_entity_result(self, query_entity_result):
        """
        处理实体抽取结果，包括时间实体替换逻辑和normal_entity设置

        Args:
            query_entity_result: 实体抽取结果字典
        """
        try:
            # 处理normal_entity，设置到entity_slot中以支持entity retrieval
            normal_entity = query_entity_result.get("normal_entity", [])
            self.entity_slot = normal_entity
            logger.info(f"Entity slot set from unified result: {self.entity_slot}")

            # 处理时间实体替换逻辑
            time_entity = query_entity_result.get("time_entity", {})
            self.query_time_replace = self.context.search_box
            time_slot_list = []

            logger.info(f"Processing time entities: {time_entity}")

            if len(time_entity) > 0:
                for key, value in time_entity.items():
                    time_slot_list.append(value)

            from utils.date_convert import DateTimeUtil
            time_change_slot_list = DateTimeUtil.run(time_slot_list)
            time_change_slot_dict = {}
            for item in time_change_slot_list:
                entity_item = item.split("=")
                if len(entity_item) == 2:
                    time_change_slot_dict[entity_item[0]] = entity_item[1]

            slot_replace = {}
            if len(time_entity) > 0:
                for key, value in time_entity.items():
                    if value in time_change_slot_dict:
                        slot_replace[key] = time_change_slot_dict[value]

            logger.info(f"Time slot replacements: {slot_replace}")
            for key, value in slot_replace.items():
                self.query_time_replace = self.query_time_replace.replace(key, value, 1)
            logger.info(f"Original query: {self.context.search_box}")
            logger.info(f"Time-replaced query: {self.query_time_replace}")

        except Exception as e:
            logger.error(f"Error processing entity result: {e}")
            self.query_time_replace = self.context.search_box
            # 确保即使出错也设置空的entity_slot
            self.entity_slot = []

    def _process_rewrite_result(self, query_rewrite_result):
        """
        处理问题改写结果

        Args:
            query_rewrite_result: 问题改写结果字典
        """
        try:
            query_rewrite_intent = query_rewrite_result.get("intent")
            self.context.query_rewrite = query_rewrite_result.get("query")

            if query_rewrite_intent == "ask_in_reply":
                self._set_ask_in_reply_result()
            else:
                self.answer.query_rewrite = query_rewrite_result.get("query")
                self.answer.query = self.context.search_box
                self.transition(QueryState.INTENT_RECOGNITION)

        except Exception as e:
            logger.error(f"Error processing rewrite result: {e}")
            self.context.query_rewrite = self.query_time_replace
            self.transition(QueryState.INTENT_RECOGNITION)

    def _process_intent_result(self, intent_result):
        """
        处理意图分类结果并进行状态转换

        Args:
            intent_result: 意图分类结果字典
        """
        try:
            classification_intent = intent_result.get("intent", "normal_search")
            logger.info(f"Processing classification_intent: {classification_intent}")

            # 重置所有意图标志位
            self.reject_intent_flag = False
            self.knowledge_search_flag = False
            self.agent_intent_flag = False
            self.search_intent_flag = False

            # 设置对应的意图标志位（保持与_process_intent_response的逻辑一致）
            if classification_intent == "reject_search":
                self.reject_intent_flag = True
                self.search_intent_flag = False
                logger.info("Set reject_intent_flag = True, search_intent_flag = False")
            elif classification_intent == "agent_search":
                self.agent_intent_flag = True
                # 根据agent_cot_flag决定search_intent_flag
                if self.context.agent_cot_flag:
                    self.search_intent_flag = False
                else:
                    self.search_intent_flag = True
                    self.agent_intent_flag = False
                logger.info(
                    f"Set agent_intent_flag = {self.agent_intent_flag}, search_intent_flag = {self.search_intent_flag}")
            elif classification_intent == "knowledge_search":
                self.knowledge_search_flag = True
                self.search_intent_flag = False
                self.agent_intent_flag = False
                logger.info("Set knowledge_search_flag = True, search_intent_flag = False")
            else:
                self.search_intent_flag = True
                logger.info("Default to normal_search intent, search_intent_flag = True")

            # 进行状态转换
            self._transition_based_on_intent()

        except Exception as e:
            logger.error(f"Error processing intent result: {e}")
            # 默认转换到意图识别状态
            self.transition(QueryState.INTENT_RECOGNITION)

    def _handle_query_rewrite(self):
        try:
            query_rewrite_result, model_response = get_query_rewrite(
                self.context.model_type,
                self.query_time_replace,
                self.context.database_profile["prompt_map"],
                self.context.user_query_history,
                self.context.database_profile["prompt_environment"],
            )
            self.token_info[QueryState.QUERY_REWRITE.name] = model_response.token_info
            query_rewrite_intent = query_rewrite_result.get("intent")
            self.context.query_rewrite = query_rewrite_result.get("query")
            if query_rewrite_intent == "ask_in_reply":
                self._set_ask_in_reply_result()
            else:
                self.answer.query_rewrite = query_rewrite_result.get("query")
                self.answer.query = self.context.search_box
                self.transition(QueryState.INTENT_RECOGNITION)
        except Exception as e:
            traceback.print_exc()
            self.answer.error_log[QueryState.QUERY_REWRITE.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, handle_initial encountered an error: {e}"
            )
            self.transition(QueryState.ERROR)

    def _set_ask_in_reply_result(self):
        self.answer.query = self.context.search_box
        self.answer.query_intent = "ask_in_reply"
        self.answer.query_rewrite = self.context.query_rewrite

        self.answer.ask_rewrite_result.query_rewrite = self.context.query_rewrite
        self.transition(QueryState.COMPLETE)

    def handle_entity_retrieval(self):
        logger.info(f"Starting entity retrieval with entity_slot: {self.entity_slot}")
        try:
            self.normal_search_entity_slot = self._perform_entity_retrieval()
            logger.info(f"Entity retrieval completed, found {len(self.normal_search_entity_slot)} entities")
            entity_retrieve = []
            same_name_entity = {}
            for each_entity in self.normal_search_entity_slot:
                each_item_dict = {}
                each_item_dict["_score"] = each_entity["_score"]
                each_item_dict["_source"] = each_entity["_source"]
                if "vector_field" in each_item_dict["_source"]:
                    del each_item_dict["_source"]["vector_field"]
                entity_retrieve.append(each_item_dict)
                if (
                        each_entity["_source"]["entity_count"] > 1
                        and each_entity["_score"] > 0.98
                ):
                    same_name_entity[each_entity["_source"]["entity"]] = each_entity[
                        "_source"
                    ]["entity_table_info"]
            if (
                    len(same_name_entity) > 0
                    and self.answer.query_intent == "normal_search"
            ):
                for key, value in same_name_entity.items():
                    change_value = []
                    for each_value in value:
                        new_each_value = each_value
                        new_each_value["id"] = (
                                new_each_value["table_name"]
                                + "#"
                                + new_each_value["column_name"]
                                + "#"
                                + new_each_value["value"]
                        )
                        new_each_value["text"] = (
                                "实体名称："
                                + key
                                + ", 数据表："
                                + each_value["table_name"]
                                + "，"
                                + "列名是："
                                + each_value["column_name"]
                                + "，"
                                + "查询值是："
                                + each_value["value"]
                                + "\n"
                        )
                        change_value.append(new_each_value)
                    same_name_entity[key] = change_value

                self.answer.ask_entity_select.entity_select_info = same_name_entity
                self.answer.ask_entity_select.entity_retrieval = entity_retrieve
                self.transition(QueryState.ASK_ENTITY_SELECT)
            else:
                self.transition(QueryState.QA_RETRIEVAL)
        except Exception as e:
            traceback.print_exc()
            self.answer.error_log[QueryState.ENTITY_RETRIEVAL.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, handle_entity_retrieval encountered an error: {e}"
            )
            self.transition(QueryState.ERROR)

    def _perform_entity_retrieval(self):
        if self.context.use_rag_flag:
            return entity_retrieve_search_optimized(
                self.entity_slot,
                self.context.opensearch_info,
                self.context.selected_profile,
                parallel_threshold=2,
                max_workers=5
            )
        else:
            return []

    @log_execution
    def handle_qa_retrieval(self):
        try:
            self.normal_search_qa_retrival = self._perform_qa_retrieval()
            self.transition(QueryState.SQL_GENERATION)
        except Exception as e:
            self.answer.error_log[QueryState.QA_RETRIEVAL.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, handle_qa_retrieval encountered an error: {e}"
            )
            self.transition(QueryState.ERROR)

    def _perform_qa_retrieval(self):
        if self.context.use_rag_flag:
            return qa_retrieve_search(
                self.context.query_rewrite,
                self.context.opensearch_info,
                self.context.selected_profile,
            )
        return []

    @log_execution
    def handle_sql_generation(self):
        sql, response, original_sql = self._generate_sql()
        self.intent_search_result["sql"] = sql
        self.intent_search_result["response"] = response
        self.intent_search_result["original_sql"] = original_sql
        self.answer.sql_search_result.sql = sql.strip()
        self.answer.sql_search_result.sql_gen_process = get_generated_sql_explain(
            response
        ).strip()
        if self.context.visualize_results_flag:
            self.transition(QueryState.EXECUTE_QUERY)
        else:
            self.transition(QueryState.COMPLETE)

    def _apply_row_level_security_for_sql(self, sql):
        post_sql = DataSourceFactory.apply_row_level_security_for_sql(
            self.context.database_profile["db_type"],
            sql,
            self.context.database_profile["row_level_security_config"],
            self.context.username,
        )
        return post_sql

    def _generate_sql(self):
        try:
            response, model_response = text_to_sql(
                self.context.database_profile["tables_info"],
                self.context.database_profile["hints"],
                self.context.database_profile["prompt_map"],
                self.context.query_rewrite,
                model_id=self.context.model_type,
                sql_examples=self.normal_search_qa_retrival,
                ner_example=self.normal_search_entity_slot,
                dialect=self.context.database_profile["db_type"],
                environment_dict=self.context.database_profile["prompt_environment"],
                limit=self.context.limit,
            )
            self.token_info[QueryState.SQL_GENERATION.name] = model_response.token_info
            sql = get_generated_sql(response)
            # post-processing the sql
            # post_sql = self._apply_row_level_security_for_sql(sql)
            return sql, response, sql
        except Exception as e:
            traceback.print_exc()
            self.answer.error_log[QueryState.SQL_GENERATION.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, _generate_sql encountered an error: {e}"
            )
            return "", "", ""

    def _generate_sql_again(self):
        try:
            response, model_response = text_to_sql(
                self.context.database_profile["tables_info"],
                self.context.database_profile["hints"],
                self.context.database_profile["prompt_map"],
                self.context.query_rewrite,
                model_id=self.context.model_type,
                sql_examples=self.normal_search_qa_retrival,
                ner_example=self.normal_search_entity_slot,
                dialect=self.context.database_profile["db_type"],
                model_provider=None,
                environment_dict=self.context.database_profile["prompt_environment"],
                limit=self.context.limit,
                additional_info="""\n NOTE: when I try to write a SQL <sql>{sql_statement}</sql>, I got an error <error>{error}</error>. Please consider and avoid this problem. """.format(
                    sql_statement=self.intent_search_result["original_sql"],
                    error=self.intent_search_result["sql_execute_result"]["error_info"],
                ),
            )
            self.token_info[QueryState.SQL_GENERATION.name + "_AGAIN"] = (
                model_response.token_info
            )
            sql = get_generated_sql(response)
            # post_sql = self._apply_row_level_security_for_sql(sql) # 默认不要这一步，因为在其他地方接入了鉴权
            self.delete_error_log_entry(QueryState.SQL_GENERATION.name)
            return sql, response, sql
        except Exception as e:
            traceback.print_exc()
            self.answer.error_log[QueryState.SQL_GENERATION.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, _generate_sql encountered an error: {e}"
            )
            return "", "", ""

    @log_execution
    def handle_agent_sql_generation(self):
        agent_search_result, token_info = agent_text_search(
            self.context.query_rewrite,
            self.context.model_type,
            self.context.database_profile,
            self.entity_slot,
            self.context.opensearch_info,
            self.context.selected_profile,
            self.context.use_rag_flag,
            self.agent_task_split,
            environment_dict=self.context.database_profile["prompt_environment"],
            limit=self.context.limit,
        )
        self.token_info[QueryState.SQL_GENERATION.name + "AGENT"] = token_info
        self.agent_search_result = agent_search_result
        self.transition(QueryState.AGENT_DATA_SUMMARY)

    async def handle_agent_sql_generation_websocket(
            self, websocket, session_id, user_id
    ):
        state_name = self.get_state().name if self.get_state() else "Unknown State"
        logger.info(
            f"Executing handle_agent_sql_generation_websocket in state {state_name}"
        )
        agent_search_result, token_info = await agent_text_search_websocket(
            websocket,
            session_id,
            user_id,
            self.context.query_rewrite,
            self.context.model_type,
            self.context.database_profile,
            self.entity_slot,
            self.context.opensearch_info,
            self.context.selected_profile,
            self.context.use_rag_flag,
            self.agent_task_split,
            environment_dict=self.context.database_profile["prompt_environment"],
            limit=self.context.limit,
        )
        self.token_info[QueryState.SQL_GENERATION.name + "AGENT"] = token_info
        self.agent_search_result = agent_search_result
        self.transition(QueryState.AGENT_DATA_SUMMARY)

    @log_execution
    def handle_intent_recognition(self):
        try:
            if self.context.intent_ner_recognition_flag:
                intent_response, model_response = get_query_intent(
                    self.context.model_type,
                    self.context.query_rewrite,
                    self.context.database_profile["prompt_map"],
                    self.context.database_profile["prompt_environment"],
                )
                self.token_info[QueryState.INTENT_RECOGNITION.name] = (
                    model_response.token_info
                )
                self.intent_response = intent_response
                self._process_intent_response(intent_response)
            else:
                self.search_intent_flag = True
            self._transition_based_on_intent()
        except Exception as e:
            traceback.print_exc()
            self.answer.error_log[QueryState.INTENT_RECOGNITION.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, handle_intent_recognition encountered an error: {e}"
            )
            self.transition(QueryState.ERROR)

    def _process_intent_response(self, intent_response):
        intent = intent_response.get("intent", "normal_search")
        self.entity_slot = intent_response.get("slot", [])
        if intent == "reject_search":
            self.reject_intent_flag = True
            self.search_intent_flag = False
        elif intent == "agent_search":
            self.agent_intent_flag = True
            if self.context.agent_cot_flag:
                self.search_intent_flag = False
            else:
                self.search_intent_flag = True
                self.agent_intent_flag = False
        elif intent == "knowledge_search":
            self.knowledge_search_flag = True
            self.search_intent_flag = False
            self.agent_intent_flag = False
        else:
            self.search_intent_flag = True

    def _transition_based_on_intent(self):
        logger.info(
            f"Intent transition check - reject_intent_flag: {self.reject_intent_flag}, knowledge_search_flag: {self.knowledge_search_flag}, agent_intent_flag: {self.agent_intent_flag}")

        if self.reject_intent_flag:
            logger.info("Transitioning to REJECT_INTENT state")
            self.answer.query_intent = "reject_search"
            self.transition(QueryState.REJECT_INTENT)
        elif self.knowledge_search_flag:
            logger.info("Transitioning to KNOWLEDGE_SEARCH state")
            self.answer.query_intent = "knowledge_search"
            self.transition(QueryState.KNOWLEDGE_SEARCH)
        elif self.agent_intent_flag:
            logger.info("Transitioning to AGENT_TASK state")
            self.answer.query_intent = "agent_search"
            self.transition(QueryState.AGENT_TASK)
        else:
            logger.info(f"Transitioning to ENTITY_RETRIEVAL state with entity_slot: {self.entity_slot}")
            self.answer.query_intent = "normal_search"
            self.transition(QueryState.ENTITY_RETRIEVAL)

    @log_execution
    def handle_reject_intent(self):
        self.answer.query = self.context.search_box
        self.answer.query_rewrite = self.context.query_rewrite
        self.answer.query_intent = "reject_search"
        self.transition(QueryState.COMPLETE)

    @log_execution
    def handle_knowledge_search(self):
        response, model_response = knowledge_search(
            search_box=self.context.query_rewrite,
            ddl=self.context.database_profile["tables_info"],
            model_id=self.context.model_type,
            prompt_map=self.context.database_profile["prompt_map"],
            environment_dict=self.context.database_profile["prompt_environment"],
        )
        self.token_info[QueryState.KNOWLEDGE_SEARCH.name] = model_response.token_info
        self.answer.query = self.context.search_box
        self.answer.query_rewrite = self.context.query_rewrite
        self.answer.query_intent = "knowledge_search"
        self.answer.knowledge_search_result.knowledge_response = response
        self.transition(QueryState.COMPLETE)

    @log_execution
    def handle_entity_selection(self):
        # Handle entity selection
        entity_select_format = (
            "根据您的描述，检索到多个相同名称的实体，请选择您想要查询的实体。\n"
        )
        alphabet_list = list(string.ascii_uppercase)
        index = 0
        for entity in self.answer.ask_entity_select.entity_select_info:
            entity_value = self.answer.ask_entity_select.entity_select_info[entity]
            entity_name = entity
            entity_desc = "实体：" + entity_name + "，有如下维度值：\n"
            for each_value in entity_value:
                if index < len(alphabet_list):
                    entity_desc += alphabet_list[index] + " ："
                    entity_desc = (
                            entity_desc
                            + "数据表："
                            + each_value["table_name"]
                            + "，"
                            + "列名是："
                            + each_value["column_name"]
                            + "，"
                            + "查询值是："
                            + each_value["value"]
                            + "\n"
                    )
                    index = index + 1
            entity_select_format += entity_desc
        self.answer.query_intent = "entity_select"
        self.transition(QueryState.COMPLETE)

    @log_execution
    def handle_execute_query(self):
        try:
            sql = self.intent_search_result.get("sql", "")
            sql_execute_result = self._execute_sql(sql)
            self.intent_search_result["sql_execute_result"] = sql_execute_result
            self.answer.sql_search_result.sql_data = sql_execute_result["data"]
            self.answer.sql_search_result.sql_data_url = sql_execute_result["data_url"]
            if (
                    self.context.data_with_analyse
                    and sql_execute_result["status_code"] == 200
            ):
                self.transition(QueryState.ANALYZE_DATA)
            elif sql_execute_result["status_code"] == 200:
                self.transition(QueryState.COMPLETE)
            elif (
                    sql_execute_result["status_code"] == 500
                    and self.context.auto_correction_flag
            ):
                self.use_auto_correction_flag = True
                self.first_sql_execute_info = sql_execute_result
                sql, response, original_sql = self._generate_sql_again()
                sql_execute_result = self._execute_sql(sql)
                self.answer.sql_search_result.sql = sql
                self.answer.sql_search_result.sql_data_url = sql_execute_result[
                    "data_url"
                ]
                self.answer.sql_search_result.sql_gen_process = (
                    get_generated_sql_explain(response)
                )
                self.intent_search_result["sql_execute_result"] = sql_execute_result
                self.answer.sql_search_result.sql_data = sql_execute_result["data"]
                if (
                        self.context.data_with_analyse
                        and sql_execute_result["status_code"] == 200
                ):
                    self.transition(QueryState.ANALYZE_DATA)
                elif sql_execute_result["status_code"] == 200:
                    self.transition(QueryState.COMPLETE)
                else:
                    self.answer.error_log[QueryState.EXECUTE_QUERY.name] = (
                        sql_execute_result["error_info"]
                    )
                    self.transition(QueryState.ERROR)
            else:
                self.answer.error_log[QueryState.EXECUTE_QUERY.name] = (
                    sql_execute_result["error_info"]
                )
                self.transition(QueryState.ERROR)
        except Exception as e:
            traceback.print_exc()
            self.answer.error_log[QueryState.EXECUTE_QUERY.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, handle_execute_query encountered an error: {e}"
            )
            self.transition(QueryState.ERROR)

    def _execute_sql(self, sql):
        if sql == "":
            return {
                "data": [],
                "data_url": "",
                "sql": sql,
                "status_code": 500,
                "error_info": "The SQL is empty.",
            }
        return get_sql_result_tool(self.context.database_profile, sql)

    @log_execution
    def handle_analyze_data(self):
        # Analyze the data
        try:
            normal_search_entity_info = {}
            for item in self.normal_search_entity_slot:
                normal_search_entity_info[item["_source"]["entity"]] = item["_source"][
                    "comment"
                ]
            analyze_data = self.intent_search_result["sql_execute_result"]["data"]
            analyze_data = analyze_data.loc[:, ~analyze_data.columns.duplicated()]
            search_intent_analyse_result, model_response = data_analyse_tool(
                self.context.model_type,
                self.context.database_profile["prompt_map"],
                self.context.query_rewrite,
                analyze_data.to_json(orient="records", force_ascii=False),
                "query",
                self.context.database_profile["prompt_environment"]
            )
            if search_intent_analyse_result == "-1":
                self.answer.error_log[QueryState.ANALYZE_DATA.name] = str(
                    model_response
                )
                self.answer.sql_search_result.data_analyse = ""
                self.transition(QueryState.COMPLETE)
            else:
                self.token_info[QueryState.ANALYZE_DATA.name] = (
                    model_response.token_info
                )
                self.answer.sql_search_result.data_analyse = (
                    search_intent_analyse_result
                )
                self.transition(QueryState.COMPLETE)
        except Exception as e:
            traceback.print_exc()
            self.answer.error_log[QueryState.ANALYZE_DATA.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, handle_analyze_data encountered an error: {e}"
            )
            self.transition(QueryState.ERROR)

    @log_execution
    def handle_agent_task(self):
        # Analyze the task
        try:
            self.agent_cot_retrieve = get_retrieve_opensearch(
                self.context.opensearch_info,
                self.context.query_rewrite,
                "agent",
                self.context.selected_profile,
                2,
                0.5,
            )

            agent_cot_task_result, model_response = get_agent_cot_task(
                self.context.model_type,
                self.context.database_profile["prompt_map"],
                self.context.query_rewrite,
                self.context.database_profile["tables_info"],
                self.agent_cot_retrieve,
                self.context.database_profile["prompt_environment"],
            )
            self.token_info[QueryState.AGENT_TASK.name] = model_response.token_info
            self.agent_task_split = agent_cot_task_result
            self.transition(QueryState.AGENT_SEARCH)
        except Exception as e:
            traceback.print_exc()
            self.answer.error_log[QueryState.AGENT_TASK.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, handle_agent_task encountered an error: {e}"
            )
            self.transition(QueryState.ERROR)

    @log_execution
    def handle_agent_analyze_data(self):
        # Analyze the data
        try:
            filter_deep_dive_sql_result = []
            agent_sql_search_result = []
            for i in range(len(self.agent_search_result)):
                each_task_res = get_sql_result_tool(
                    self.context.database_profile, self.agent_search_result[i]["sql"]
                )
                if (
                        each_task_res["status_code"] == 200
                        and len(each_task_res["data"]) > 0
                ):
                    self.agent_search_result[i]["data_result"] = each_task_res[
                        "data"
                    ].to_json(orient="records")
                    filter_deep_dive_sql_result.append(self.agent_search_result[i])

                    show_select_data = [
                                           list(each_task_res["data"].columns)
                                       ] + each_task_res["data"].values.tolist()
                    each_task_sql_response = get_generated_sql_explain(
                        self.agent_search_result[i]["response"]
                    )
                    sub_task_sql_result = SQLSearchResult(
                        sql_data=show_select_data,
                        sql_data_url=each_task_res["data_url"],
                        sql=self.agent_search_result[i]["sql"],
                        data_show_type="table",
                        sql_gen_process=each_task_sql_response,
                        data_analyse="",
                        sql_data_chart=[],
                    )
                    each_task_sql_search_result = TaskSQLSearchResult(
                        sub_task_query=self.agent_search_result[i]["query"],
                        sql_search_result=sub_task_sql_result,
                    )
                    agent_sql_search_result.append(each_task_sql_search_result)

            agent_data_analyse_result, model_response = data_analyse_tool(
                self.context.model_type,
                self.context.database_profile["prompt_map"],
                self.context.query_rewrite,
                json.dumps(filter_deep_dive_sql_result, ensure_ascii=False),
                "agent",
                self.context.database_profile["prompt_environment"],
            )
            self.token_info[QueryState.AGENT_DATA_SUMMARY.name] = (
                model_response.token_info
            )

            self.agent_valid_data = filter_deep_dive_sql_result
            self.agent_data_analyse_result = agent_data_analyse_result
            self.answer.agent_search_result.agent_summary = agent_data_analyse_result
            self.answer.agent_search_result.agent_sql_search_result = (
                agent_sql_search_result
            )
            self.transition(QueryState.COMPLETE)
        except Exception as e:
            traceback.print_exc()
            self.answer.error_log[QueryState.AGENT_DATA_SUMMARY.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, handle_agent_analyze_data encountered an error: {e}"
            )
            self.transition(QueryState.ERROR)

    @log_execution
    def handle_suggest_question(self):
        # Handle suggest question
        if self.context.gen_suggested_question_flag:
            if self.search_intent_flag or self.agent_intent_flag:
                generated_sq, model_response = generate_suggested_question(self.context.database_profile['prompt_map'],
                                                                           self.context.query_rewrite,
                                                                           model_id=self.context.model_type,
                                                                           environment_dict=
                                                                           self.context.database_profile[
                                                                               'prompt_environment'],
                                                                           assistant_prompt="[generate]")
                self.token_info["SUGGEST_QUESTION"] = model_response.token_info

                logger.info("generated_sq: %s", generated_sq)

                split_strings = generated_sq.split("[generate]")
                gen_sq_list = [s.strip() for s in split_strings if s.strip()]
                self.answer.suggested_question = gen_sq_list

    def delete_error_log_entry(self, key):
        if key in self.answer.error_log:
            del self.answer.error_log[key]

    def handle_data_visualization(self):
        try:
            if self.answer.query_intent == "normal_search":
                (
                    model_select_type,
                    show_select_data,
                    select_chart_type,
                    show_chart_data,
                    model_response,
                ) = data_visualization(
                    self.context.model_type,
                    self.context.query_rewrite,
                    self.get_answer().sql_search_result.sql_data,
                    self.context.database_profile["prompt_map"],
                    self.context.database_profile["prompt_environment"],
                )
                self.token_info[QueryState.DATA_VISUALIZATION.name] = (
                    model_response.token_info
                )
                if select_chart_type != "-1":
                    sql_chart_data = ChartEntity(chart_type="", chart_data=[])
                    sql_chart_data.chart_type = select_chart_type
                    sql_chart_data.chart_data = show_chart_data
                    self.get_answer().sql_search_result.sql_data_chart = [
                        sql_chart_data
                    ]
                self.get_answer().sql_search_result.data_show_type = model_select_type
                self.get_answer().sql_search_result.sql_data = show_select_data
            elif self.answer.query_intent == "agent_search":
                agent_sql_search_result = (
                    self.answer.agent_search_result.agent_sql_search_result
                )
                agent_sql_search_result_with_visualization = []
                for each in agent_sql_search_result:
                    (
                        model_select_type,
                        show_select_data,
                        select_chart_type,
                        show_chart_data,
                        model_response,
                    ) = data_visualization(
                        self.context.model_type,
                        each.sub_task_query,
                        each.sql_search_result.sql_data,
                        self.context.database_profile["prompt_map"],
                        self.context.database_profile["prompt_environment"],
                    )
                    if QueryState.DATA_VISUALIZATION.name not in self.token_info:
                        self.token_info[QueryState.DATA_VISUALIZATION.name] = (
                            model_response.token_info
                        )
                    else:
                        if "input_tokens" in model_response.token_info:
                            if (
                                    "input_tokens"
                                    not in self.token_info[
                                QueryState.DATA_VISUALIZATION.name
                            ]
                            ):
                                self.token_info[QueryState.DATA_VISUALIZATION.name][
                                    "input_tokens"
                                ] = 0
                            self.token_info[QueryState.DATA_VISUALIZATION.name][
                                "input_tokens"
                            ] = (
                                    self.token_info[QueryState.DATA_VISUALIZATION.name][
                                        "input_tokens"
                                    ]
                                    + model_response.token_info["input_tokens"]
                            )
                        if "output_tokens" in model_response.token_info:
                            if (
                                    "output_tokens"
                                    not in self.token_info[
                                QueryState.DATA_VISUALIZATION.name
                            ]
                            ):
                                self.token_info[QueryState.DATA_VISUALIZATION.name][
                                    "output_tokens"
                                ] = 0
                            self.token_info[QueryState.DATA_VISUALIZATION.name][
                                "output_tokens"
                            ] = (
                                    self.token_info[QueryState.DATA_VISUALIZATION.name][
                                        "output_tokens"
                                    ]
                                    + model_response.token_info["output_tokens"]
                            )
                    if select_chart_type != "-1":
                        sql_chart_data = ChartEntity(chart_type="", chart_data=[])
                        sql_chart_data.chart_type = select_chart_type
                        sql_chart_data.chart_data = show_chart_data
                        each.sql_search_result.sql_data_chart = [sql_chart_data]
                    each.sql_search_result.data_show_type = model_select_type
                    each.sql_search_result.sql_data = show_select_data
                    agent_sql_search_result_with_visualization.append(each)
                self.answer.agent_search_result.agent_sql_search_result = (
                    agent_sql_search_result_with_visualization
                )
        except Exception as e:
            traceback.print_exc()
            self.answer.error_log[QueryState.DATA_VISUALIZATION.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, handle_data_visualization encountered an error: {e}"
            )
            self.transition(QueryState.ERROR)

    def handle_user_select_entity(self):
        try:
            self.answer.query = self.context.search_box
            self.answer.query_rewrite = self.context.query_rewrite
            self.answer.query_intent = "normal_search"
            self.search_intent_flag = True
            comment_format = "{entity} is located in table {table_name}, column {column_name}, the dimension value is {value}"
            normal_search_entity_slot = self.context.entity_retrieval
            entity_user_select = self.context.entity_user_select
            entity_retrieval = []
            for each_entity in normal_search_entity_slot:
                entity = each_entity["_source"]["entity"]
                if entity in entity_user_select:
                    each_entity["_source"]["entity_table_info"] = [
                        entity_user_select[entity]
                    ]
                    each_entity["_source"]["comment"] = comment_format.format(
                        entity=entity,
                        table_name=entity_user_select[entity]["table_name"],
                        column_name=entity_user_select[entity]["column_name"],
                        value=entity_user_select[entity]["value"],
                    )
                    entity_retrieval.append(each_entity)
            self.normal_search_entity_slot = entity_retrieval
            logger.info(entity_retrieval)
            self.transition(QueryState.QA_RETRIEVAL)
        except Exception as e:
            traceback.print_exc()
            self.answer.error_log[QueryState.USER_SELECT_ENTITY.name] = str(e)
            logger.error(
                f"The context is {self.context.search_box}, handle_user_select encountered an error: {e}"
            )
            self.transition(QueryState.ERROR)

    def handle_add_to_log(self, log_id):
        answer_info = change_class_to_str(self.answer)
        current_time = get_current_time()
        sql = ""
        if self.answer.query_intent == "normal_search":
            sql = self.answer.sql_search_result.sql
        LogManagement.add_log_to_database(
            log_id=log_id,
            user_id=self.context.user_id,
            session_id=self.context.session_id,
            profile_name=self.context.selected_profile,
            sql=sql,
            query=self.context.search_box,
            intent=self.answer.query_intent,
            log_info=answer_info,
            log_type="chat_history",
            time_str=current_time,
        )

    def answer_data_check(self):
        # Check and convert sql_data if it's a DataFrame
        if isinstance(self.answer.sql_search_result.sql_data, pd.DataFrame):
            # Convert DataFrame to list format
            columns = list(self.answer.sql_search_result.sql_data.columns)
            rows = self.answer.sql_search_result.sql_data.values.tolist()
            self.answer.sql_search_result.sql_data = [columns] + rows